import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import (QApplication, QMainWindow, QTabWidget, QWidget, QVBoxLayout,
                             QHBoxLayout, QLabel, QPushButton, QFileDialog, QSlider,
                             QComboBox, QMessageBox, QGroupBox, QGridLayout, QAction,
                             QToolBar, QSplitter, QDockWidget, QTextEdit, QProgressBar)
from PyQt5.QtGui import QPixmap, QImage, QIcon, QFont
from PyQt5.QtCore import Qt, QTimer, QThread, pyqtSignal
from PyQt5 import QtCore


# 视频处理线程类
class VideoThread(QThread):
    change_pixmap_signal = pyqtSignal(np.ndarray)
    finished_signal = pyqtSignal()

    def __init__(self, video_path=None):
        super().__init__()
        self.video_path = video_path
        self._run_flag = True
        self.processing_function = lambda frame: frame  # 默认不处理

    def run(self):
        # 打开摄像头或视频文件
        if self.video_path:
            cap = cv2.VideoCapture(self.video_path)
        else:
            cap = cv2.VideoCapture(0)

        while self._run_flag:
            ret, cv_img = cap.read()
            if ret:
                # 应用图像处理函数
                processed_img = self.processing_function(cv_img)
                self.change_pixmap_signal.emit(processed_img)
            else:
                break
            # 控制帧率
            self.msleep(30)

        # 释放资源
        cap.release()
        self.finished_signal.emit()

    def stop(self):
        """停止线程"""
        self._run_flag = False
        self.wait()


# 主窗口类
class ImageVideoProcessor(QMainWindow):
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.current_image = None
        self.current_video_thread = None
        self.is_video_playing = False
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

    def init_ui(self):
        # 设置窗口
        self.setWindowTitle('图像处理与视频分析平台')
        self.setGeometry(100, 100, 1200, 800)

        # 创建菜单栏
        self.create_menu_bar()

        # 创建工具栏
        self.create_tool_bar()

        # 创建中央部件
        self.central_widget = QTabWidget()
        self.setCentralWidget(self.central_widget)

        # 添加图像和视频处理选项卡
        self.image_tab = QWidget()
        self.video_tab = QWidget()
        self.central_widget.addTab(self.image_tab, '图像处理')
        self.central_widget.addTab(self.video_tab, '视频处理')

        # 设置图像选项卡
        self.setup_image_tab()

        # 设置视频选项卡
        self.setup_video_tab()

        # 创建状态栏
        self.statusBar().showMessage('就绪')

    def create_menu_bar(self):
        # 创建菜单栏
        menubar = self.menuBar()

        # 文件菜单
        file_menu = menubar.addMenu('文件')

        # 打开图像动作
        open_image_action = QAction('打开图像', self)
        open_image_action.setShortcut('Ctrl+I')
        open_image_action.triggered.connect(self.open_image)
        file_menu.addAction(open_image_action)

        # 打开视频动作
        open_video_action = QAction('打开视频', self)
        open_video_action.setShortcut('Ctrl+V')
        open_video_action.triggered.connect(self.open_video)
        file_menu.addAction(open_video_action)

        # 保存图像动作
        save_image_action = QAction('保存图像', self)
        save_image_action.setShortcut('Ctrl+S')
        save_image_action.triggered.connect(self.save_image)
        file_menu.addAction(save_image_action)

        file_menu.addSeparator()

        # 退出动作
        exit_action = QAction('退出', self)
        exit_action.setShortcut('Ctrl+Q')
        exit_action.triggered.connect(self.close)
        file_menu.addAction(exit_action)

        # 帮助菜单
        help_menu = menubar.addMenu('帮助')

        # 关于动作
        about_action = QAction('关于', self)
        about_action.triggered.connect(self.show_about)
        help_menu.addAction(about_action)

    def create_tool_bar(self):
        # 创建工具栏
        toolbar = QToolBar('工具栏')
        self.addToolBar(toolbar)

        # 添加图像工具按钮
        open_image_btn = QPushButton('打开图像')
        open_image_btn.clicked.connect(self.open_image)
        toolbar.addWidget(open_image_btn)

        save_image_btn = QPushButton('保存图像')
        save_image_btn.clicked.connect(self.save_image)
        toolbar.addWidget(save_image_btn)

        toolbar.addSeparator()

        # 添加视频工具按钮
        open_video_btn = QPushButton('打开视频')
        open_video_btn.clicked.connect(self.open_video)
        toolbar.addWidget(open_video_btn)

        start_video_btn = QPushButton('开始播放')
        start_video_btn.clicked.connect(self.start_video)
        toolbar.addWidget(start_video_btn)

        stop_video_btn = QPushButton('停止播放')
        stop_video_btn.clicked.connect(self.stop_video)
        toolbar.addWidget(stop_video_btn)

    def setup_image_tab(self):
        # 创建主布局
        main_layout = QHBoxLayout(self.image_tab)

        # 创建左侧控制面板
        control_panel = QWidget()
        control_layout = QVBoxLayout(control_panel)
        control_panel.setMaximumWidth(300)

        # 图像加载组
        load_group = QGroupBox('图像加载')
        load_layout = QVBoxLayout(load_group)

        load_btn = QPushButton('加载图像')
        load_btn.clicked.connect(self.open_image)
        load_layout.addWidget(load_btn)

        control_layout.addWidget(load_group)

        # 图像处理组
        process_group = QGroupBox('图像处理')
        process_layout = QGridLayout(process_group)

        # 灰度转换按钮
        gray_btn = QPushButton('灰度转换')
        gray_btn.clicked.connect(lambda: self.process_image('gray'))
        process_layout.addWidget(gray_btn, 0, 0)

        # 边缘检测按钮
        edge_btn = QPushButton('边缘检测')
        edge_btn.clicked.connect(lambda: self.process_image('edge'))
        process_layout.addWidget(edge_btn, 0, 1)

        # 模糊处理按钮
        blur_btn = QPushButton('模糊处理')
        blur_btn.clicked.connect(lambda: self.process_image('blur'))
        process_layout.addWidget(blur_btn, 1, 0)

        # 锐化处理按钮
        sharpen_btn = QPushButton('锐化处理')
        sharpen_btn.clicked.connect(lambda: self.process_image('sharpen'))
        process_layout.addWidget(sharpen_btn, 1, 1)

        # 二值化按钮
        thresh_btn = QPushButton('二值化')
        thresh_layout = QHBoxLayout()
        self.thresh_value = QSlider(Qt.Horizontal)
        self.thresh_value.setRange(0, 255)
        self.thresh_value.setValue(127)
        thresh_layout.addWidget(self.thresh_value)
        thresh_btn.clicked.connect(lambda: self.process_image('threshold'))
        thresh_layout.addWidget(thresh_btn)
        process_layout.addLayout(thresh_layout, 2, 0, 1, 2)

        control_layout.addWidget(process_group)

        # 形状识别组
        shape_group = QGroupBox('形状识别')
        shape_layout = QVBoxLayout(shape_group)

        # 人脸检测按钮
        face_btn = QPushButton('人脸检测')
        face_btn.clicked.connect(lambda: self.process_image('face'))
        shape_layout.addWidget(face_btn)

        # 圆形检测按钮
        circle_btn = QPushButton('圆形检测')
        circle_btn.clicked.connect(lambda: self.process_image('circle'))
        shape_layout.addWidget(circle_btn)

        # 矩形检测按钮
        rectangle_btn = QPushButton('矩形检测')
        rectangle_btn.clicked.connect(lambda: self.process_image('rectangle'))
        shape_layout.addWidget(rectangle_btn)

        control_layout.addWidget(shape_group)

        # 图像拼接组
        stitch_group = QGroupBox('图像拼接')
        stitch_layout = QVBoxLayout(stitch_group)

        stitch_btn = QPushButton('添加到拼接')
        stitch_btn.clicked.connect(self.add_to_stitch)
        stitch_layout.addWidget(stitch_btn)

        stitch_process_btn = QPushButton('执行拼接')
        stitch_process_btn.clicked.connect(self.stitch_images)
        stitch_layout.addWidget(stitch_process_btn)

        control_layout.addWidget(stitch_group)

        # 添加到主布局
        main_layout.addWidget(control_panel)

        # 创建右侧图像显示区域
        self.image_display = QLabel('请加载图像')
        self.image_display.setAlignment(Qt.AlignCenter)
        self.image_display.setMinimumSize(400, 400)
        main_layout.addWidget(self.image_display, 1)

        # 初始化图像拼接列表
        self.images_to_stitch = []

    def setup_video_tab(self):
        # 创建主布局
        main_layout = QVBoxLayout(self.video_tab)

        # 创建视频控制区域
        control_layout = QHBoxLayout()

        # 打开视频按钮
        open_video_btn = QPushButton('打开视频')
        open_video_btn.clicked.connect(self.open_video)
        control_layout.addWidget(open_video_btn)

        # 开始播放按钮
        self.start_video_btn = QPushButton('开始播放')
        self.start_video_btn.clicked.connect(self.start_video)
        control_layout.addWidget(self.start_video_btn)

        # 停止播放按钮
        self.stop_video_btn = QPushButton('停止播放')
        self.stop_video_btn.clicked.connect(self.stop_video)
        self.stop_video_btn.setEnabled(False)
        control_layout.addWidget(self.stop_video_btn)

        # 选择摄像头按钮
        self.camera_btn = QPushButton('使用摄像头')
        self.camera_btn.clicked.connect(self.use_camera)
        control_layout.addWidget(self.camera_btn)

        # 添加视频处理选择
        self.video_process_combo = QComboBox()
        self.video_process_combo.addItems(['无处理', '灰度', '边缘检测', '人脸检测'])
        self.video_process_combo.currentIndexChanged.connect(self.change_video_processing)
        control_layout.addWidget(self.video_process_combo)

        main_layout.addLayout(control_layout)

        # 创建视频显示区域
        self.video_display = QLabel('请打开视频或选择摄像头')
        self.video_display.setAlignment(Qt.AlignCenter)
        self.video_display.setMinimumSize(640, 480)
        main_layout.addWidget(self.video_display, 1)

        # 创建进度条
        self.video_progress = QProgressBar()
        self.video_progress.setRange(0, 100)
        main_layout.addWidget(self.video_progress)

    def open_image(self):
        options = QFileDialog.Options()
        file_name, _ = QFileDialog.getOpenFileName(
            self, "打开图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp *.gif);;所有文件 (*)", options=options)

        if file_name:
            self.current_image = cv2.imread(file_name)
            self.display_image(self.current_image)
            self.statusBar().showMessage(f'加载图像: {file_name}')

    def display_image(self, img):
        if img is not None:
            # 转换为RGB
            if len(img.shape) == 3:
                rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            else:
                rgb_img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)

            # 创建QImage
            h, w, ch = rgb_img.shape
            bytes_per_line = ch * w
            q_img = QImage(rgb_img.data, w, h, bytes_per_line, QImage.Format_RGB888)

            # 显示图像
            self.image_display.setPixmap(QPixmap.fromImage(q_img).scaled(
                self.image_display.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

    def process_image(self, operation):
        if self.current_image is None:
            QMessageBox.warning(self, '警告', '请先加载图像')
            return

        processed_img = None

        if operation == 'gray':
            processed_img = cv2.cvtColor(self.current_image, cv2.COLOR_BGR2GRAY)

        elif operation == 'edge':
            gray = cv2.cvtColor(self.current_image, cv2.COLOR_BGR2GRAY)
            processed_img = cv2.Canny(gray, 100, 200)

        elif operation == 'blur':
            processed_img = cv2.GaussianBlur(self.current_image, (15, 15), 0)

        elif operation == 'sharpen':
            kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
            processed_img = cv2.filter2D(self.current_image, -1, kernel)

        elif operation == 'threshold':
            gray = cv2.cvtColor(self.current_image, cv2.COLOR_BGR2GRAY)
            thresh_value = self.thresh_value.value()
            _, processed_img = cv2.threshold(gray, thresh_value, 255, cv2.THRESH_BINARY)

        elif operation == 'face':
            processed_img = self.current_image.copy()
            gray = cv2.cvtColor(processed_img, cv2.COLOR_BGR2GRAY)
            faces = self.face_cascade.detectMultiScale(gray, 1.1, 4)

            for (x, y, w, h) in faces:
                cv2.rectangle(processed_img, (x, y), (x + w, y + h), (255, 0, 0), 2)

            self.statusBar().showMessage(f'检测到 {len(faces)} 个人脸')

        elif operation == 'circle':
            processed_img = self.current_image.copy()
            gray = cv2.cvtColor(processed_img, cv2.COLOR_BGR2GRAY)
            gray = cv2.medianBlur(gray, 5)

            circles = cv2.HoughCircles(gray, cv2.HOUGH_GRADIENT, 1, 20,
                                       param1=50, param2=30, minRadius=0, maxRadius=0)

            if circles is not None:
                circles = np.uint16(np.around(circles))
                for i in circles[0, :]:
                    # 绘制外圆
                    cv2.circle(processed_img, (i[0], i[1]), i[2], (0, 255, 0), 2)
                    # 绘制圆心
                    cv2.circle(processed_img, (i[0], i[1]), 2, (0, 0, 255), 3)

                self.statusBar().showMessage(f'检测到 {len(circles[0])} 个圆形')

        elif operation == 'rectangle':
            processed_img = self.current_image.copy()
            gray = cv2.cvtColor(processed_img, cv2.COLOR_BGR2GRAY)
            _, thresh = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)

            contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

            rectangles = []
            for contour in contours:
                epsilon = 0.02 * cv2.arcLength(contour, True)
                approx = cv2.approxPolyDP(contour, epsilon, True)

                if len(approx) == 4:
                    rectangles.append(approx)
                    cv2.drawContours(processed_img, [approx], 0, (0, 0, 255), 2)

            self.statusBar().showMessage(f'检测到 {len(rectangles)} 个矩形')

        if processed_img is not None:
            self.current_image = processed_img
            self.display_image(processed_img)

    def add_to_stitch(self):
        if self.current_image is not None:
            self.images_to_stitch.append(self.current_image)
            self.statusBar().showMessage(f'已添加图像到拼接列表，当前列表大小: {len(self.images_to_stitch)}')
        else:
            QMessageBox.warning(self, '警告', '请先加载图像')

    def stitch_images(self):
        if len(self.images_to_stitch) < 2:
            QMessageBox.warning(self, '警告', '至少需要两张图像进行拼接')
            return

        stitcher = cv2.Stitcher_create()
        status, stitched_img = stitcher.stitch(self.images_to_stitch)

        if status == cv2.Stitcher_OK:
            self.current_image = stitched_img
            self.display_image(stitched_img)
            self.statusBar().showMessage('图像拼接成功')
        else:
            QMessageBox.warning(self, '拼接失败', f'拼接失败，错误代码: {status}')
            self.statusBar().showMessage('图像拼接失败')

    def save_image(self):
        if self.current_image is None:
            QMessageBox.warning(self, '警告', '请先加载或处理图像')
            return

        options = QFileDialog.Options()
        file_name, _ = QFileDialog.getSaveFileName(
            self, "保存图像", "", "PNG文件 (*.png);;JPEG文件 (*.jpg);;所有文件 (*)", options=options)

        if file_name:
            cv2.imwrite(file_name, self.current_image)
            self.statusBar().showMessage(f'图像已保存到: {file_name}')

    def open_video(self):
        options = QFileDialog.Options()
        file_name, _ = QFileDialog.getOpenFileName(
            self, "打开视频", "", "视频文件 (*.mp4 *.avi *.mov *.mkv);;所有文件 (*)", options=options)

        if file_name:
            self.stop_video()
            self.current_video_path = file_name
            self.statusBar().showMessage(f'加载视频: {file_name}')

    def start_video(self):
        if hasattr(self, 'current_video_path') or self.current_video_thread is None:
            if not hasattr(self, 'current_video_path'):
                self.use_camera()
                return

            self.stop_video()

            # 创建视频线程
            self.current_video_thread = VideoThread(self.current_video_path)
            self.current_video_thread.change_pixmap_signal.connect(self.update_video_frame)
            self.current_video_thread.finished_signal.connect(self.video_finished)

            # 设置视频处理函数
            self.change_video_processing()

            # 启动线程
            self.current_video_thread.start()
            self.is_video_playing = True
            self.start_video_btn.setEnabled(False)
            self.stop_video_btn.setEnabled(True)
            self.statusBar().showMessage('视频播放中...')
        else:
            QMessageBox.warning(self, '警告', '请先打开视频')

    def stop_video(self):
        if self.current_video_thread is not None and self.is_video_playing:
            self.current_video_thread.stop()
            self.is_video_playing = False
            self.start_video_btn.setEnabled(True)
            self.stop_video_btn.setEnabled(False)
            self.statusBar().showMessage('视频已停止')

    def use_camera(self):
        self.stop_video()

        # 创建视频线程（使用摄像头）
        self.current_video_thread = VideoThread()
        self.current_video_thread.change_pixmap_signal.connect(self.update_video_frame)
        self.current_video_thread.finished_signal.connect(self.video_finished)

        # 设置视频处理函数
        self.change_video_processing()

        # 启动线程
        self.current_video_thread.start()
        self.is_video_playing = True
        self.start_video_btn.setEnabled(False)
        self.stop_video_btn.setEnabled(True)
        self.statusBar().showMessage('正在使用摄像头...')

    def update_video_frame(self, cv_img):
        # 转换为RGB
        rgb_img = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)

        # 创建QImage
        h, w, ch = rgb_img.shape
        bytes_per_line = ch * w
        q_img = QImage(rgb_img.data, w, h, bytes_per_line, QImage.Format_RGB888)

        # 显示图像
        self.video_display.setPixmap(QPixmap.fromImage(q_img).scaled(
            self.video_display.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

    def change_video_processing(self):
        if self.current_video_thread is not None:
            process_type = self.video_process_combo.currentText()

            if process_type == '无处理':
                self.current_video_thread.processing_function = lambda frame: frame
            elif process_type == '灰度':
                self.current_video_thread.processing_function = lambda frame: cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            elif process_type == '边缘检测':
                self.current_video_thread.processing_function = lambda frame: cv2.Canny(frame, 100, 200)
            elif process_type == '人脸检测':
                def detect_faces(frame):
                    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                    faces = self.face_cascade.detectMultiScale(gray, 1.1, 4)

                    for (x, y, w, h) in faces:
                        cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)

                    return frame

                self.current_video_thread.processing_function = detect_faces

    def video_finished(self):
        self.is_video_playing = False
        self.start_video_btn.setEnabled(True)
        self.stop_video_btn.setEnabled(False)
        self.statusBar().showMessage('视频播放完毕')

    def show_about(self):
        QMessageBox.about(self, '关于',
                          '图像处理与视频分析平台\n\n'
                          '版本: 1.0\n\n'
                          '功能:\n'
                          '- 图像加载与显示\n'
                          '- 图像处理（灰度、边缘检测、模糊、锐化等）\n'
                          '- 形状识别（人脸、圆形、矩形）\n'
                          '- 图像拼接\n'
                          '- 视频播放与控制\n'
                          '- 实时视频处理')

    def closeEvent(self, event):
        # 确保关闭时停止所有线程
        if self.current_video_thread is not None and self.is_video_playing:
            self.current_video_thread.stop()
        event.accept()


if __name__ == '__main__':
    # 设置中文字体
    app = QApplication(sys.argv)
    font = QFont('SimHei')
    app.setFont(font)

    # 创建并显示窗口
    window = ImageVideoProcessor()
    window.show()

    # 运行应用程序
    sys.exit(app.exec_())